import threading

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from scripts.config import load_config
from scripts.rag_manager import RAGManager
from scripts.tool_manager import Promt_Manager
from scripts.speech_manager import SpeechManager
from scripts.history_manager import HistoryManager
from scripts.extract_manager import ExtractManager
from scripts.info_need_manager import InfoNeedManager


def main():
    cfg = load_config('./config.py')

    history_manager = HistoryManager(cfg)
    info_need_manager = InfoNeedManager(cfg)
    extract_manager = ExtractManager(cfg)
    rag_manager = RAGManager(cfg)
    speech_manager = SpeechManager(cfg)
    promt_Manager = Promt_Manager(cfg)
    
    model = ChatOpenAI(
        model=cfg.get('model'),
        base_url=cfg.get('base_url'),
        api_key=cfg.get('api_key'),
        temperature=0.7
    )

    parser = StrOutputParser()

    def extract_and_store(messages):
        data = extract_manager.extract(messages)
        if data:
            rag_manager.store(data)

    
    while True:
        # from scripts.voice_manager import VoiceManager
        # voice_manager = VoiceManager(cfg)
        # user_input = voice_manager.listen()
        user_input = input()
        print(user_input)
        
        history_manager.update("user", user_input)
        messages = history_manager.get_messages()
        promt_Manager.route(user_input)

        prompts = promt_Manager.get_all()
        if prompts:
            system_prompt = "\n".join([f"{v}" for k, v in prompts.items()])
        else:
            system_prompt = "你的回复将被语音播放，请避免长篇大论，尽量用口语表达，简洁自然。"

        prompt_temp = ChatPromptTemplate.from_messages(
            [
                SystemMessage(content=system_prompt),
                MessagesPlaceholder(variable_name="messages"),
                MessagesPlaceholder(variable_name="user_input"),
            ]
        )

        chain = prompt_temp | model | parser

        need_list = info_need_manager.extract_info_needs(messages)
        extra_infos = []
        for need in need_list:
            result_str = rag_manager.search(need)
            if result_str:
                extra_infos.append(result_str)

        if extra_infos:
            supplement = "【RAG检索补充信息】\n" + "\n".join(extra_infos) + "\n"
            user_input = supplement + user_input
        
        output = ''
        for r in chain.stream(
            {
                "messages": messages[:-1],
                "user_input": [user_input],
            },
            config={'configurable':{'session_id': "xx"}}
            ):
            print(r, end='')
            output += r
        print('\n')
        speech_manager.speak(output)

        history_manager.update("assistant", output)

        messages = history_manager.get_messages()

        threading.Thread(target=extract_and_store, args=(messages,), daemon=True).start()


if __name__ == '__main__':
    main()
